#include "ForegroundMeshPipeline.h"
#include "Render/VulkanContext.h"

#include "Render/Texture/GPURenderTarget.h"
#include "Render/Material/VulkanShader.h"
#include "Render/Material/VulkanShaders.h"

void SForegroundMeshPipeline::Init()
{
    //create DescriptorSetLayout
    {
        //create descriptor set layout
        vk::DescriptorSetLayoutBinding DescriptorSetLayoutBinding[1];
        DescriptorSetLayoutBinding[0]= vk::DescriptorSetLayoutBinding()
            .setBinding(0)
            .setDescriptorType(vk::DescriptorType::eUniformBuffer)
            .setDescriptorCount(1)
            .setStageFlags(vk::ShaderStageFlagBits::eVertex)
            .setPImmutableSamplers(nullptr);

        vk::DescriptorSetLayoutCreateInfo DescriptorSetLayoutCreateInfo = vk::DescriptorSetLayoutCreateInfo()
            .setBindingCount(1)
            .setPBindings(DescriptorSetLayoutBinding);

        DescriptorSetLayout = CVulkanContext::Get()->GetDevice().createDescriptorSetLayout(DescriptorSetLayoutCreateInfo);
    }

    //add render pass
    {
        vk::AttachmentDescription AttachmentDescriptions[2];
        AttachmentDescriptions[0]= vk::AttachmentDescription()
            .setFormat(vk::Format::eR8G8B8A8Unorm)
            .setSamples(vk::SampleCountFlagBits::e1)
            .setLoadOp(vk::AttachmentLoadOp::eClear)
            .setStoreOp(vk::AttachmentStoreOp::eStore)
            .setStencilLoadOp(vk::AttachmentLoadOp::eDontCare)
            .setStencilStoreOp(vk::AttachmentStoreOp::eDontCare)
            .setInitialLayout(vk::ImageLayout::eUndefined)
            .setFinalLayout(vk::ImageLayout::eColorAttachmentOptimal);

        AttachmentDescriptions[1]= vk::AttachmentDescription()
            .setFormat(vk::Format::eD32Sfloat)
            .setSamples(vk::SampleCountFlagBits::e1)
            .setLoadOp(vk::AttachmentLoadOp::eClear)
            .setStoreOp(vk::AttachmentStoreOp::eStore)
            .setStencilLoadOp(vk::AttachmentLoadOp::eDontCare)
            .setStencilStoreOp(vk::AttachmentStoreOp::eDontCare)
            .setInitialLayout(vk::ImageLayout::eUndefined)
            .setFinalLayout(vk::ImageLayout::eDepthStencilAttachmentOptimal);

        vk::AttachmentReference ColorAttachmentReference = vk::AttachmentReference()
            .setAttachment(0)
            .setLayout(vk::ImageLayout::eColorAttachmentOptimal);

        vk::AttachmentReference DepthAttachmentReference = vk::AttachmentReference()
            .setAttachment(1)
            .setLayout(vk::ImageLayout::eDepthStencilAttachmentOptimal);

        vk::SubpassDescription SubpassDescription = vk::SubpassDescription()
            .setPipelineBindPoint(vk::PipelineBindPoint::eGraphics)
            .setColorAttachmentCount(1)
            .setPColorAttachments(&ColorAttachmentReference)
            .setPDepthStencilAttachment(&DepthAttachmentReference);

        vk::RenderPassCreateInfo RenderPassCreateInfo = vk::RenderPassCreateInfo()
            .setAttachmentCount(2)
            .setPAttachments(AttachmentDescriptions)
            .setSubpassCount(1)
            .setPSubpasses(&SubpassDescription);
        
        AxisRenderPass=CVulkanContext::Get()->GetDevice().createRenderPass(RenderPassCreateInfo);
    
    
    }

    //add axis pipeline layout
    {
        //create pipeline layout
        //pipline layout have a push constant for vertex shader
        //and vertex buffer for vertex shader
        TVector<vk::PushConstantRange> PushConstantRanges=
        {
            vk::PushConstantRange()
            .setStageFlags(vk::ShaderStageFlagBits::eVertex)
            .setOffset(0)
            .setSize(sizeof(SAxisItemVertexPushConstant))
        };


        TVector<vk::DescriptorSetLayout> DescriptorSetLayouts = {DescriptorSetLayout};

        vk::PipelineLayoutCreateInfo PipelineLayoutCreateInfo = vk::PipelineLayoutCreateInfo()
            .setSetLayouts(DescriptorSetLayouts)
            .setPushConstantRanges(PushConstantRanges);

        AxisPipelineLayout =CVulkanContext::Get()->GetDevice().createPipelineLayout(PipelineLayoutCreateInfo);

        //create pipeline
        auto VertexShader= CVulkanContext::Get()->GetShaderManager()->GetVulkanOnlyShader(U"Axis_vert");
        auto FragmentShader= CVulkanContext::Get()->GetShaderManager()->GetVulkanOnlyShader(U"Axis_frag");

        //dynamic state
        vk::DynamicState DynamicState[2]={vk::DynamicState::eViewport,vk::DynamicState::eScissor};

        vk::PipelineDynamicStateCreateInfo PipelineDynamicStateCreateInfo=vk::PipelineDynamicStateCreateInfo()
            .setDynamicStateCount(2)
            .setPDynamicStates(DynamicState);

        vk::PipelineShaderStageCreateInfo PipelineShaderStageCreateInfo[2];
        PipelineShaderStageCreateInfo[0]=vk::PipelineShaderStageCreateInfo()
            .setStage(vk::ShaderStageFlagBits::eVertex)
            .setModule(VertexShader->Module)
            .setPName("main");

        PipelineShaderStageCreateInfo[1]=vk::PipelineShaderStageCreateInfo()
            .setStage(vk::ShaderStageFlagBits::eFragment)
            .setModule(FragmentShader->Module)
            .setPName("main");

        auto VertexInputBindingDescriptions=
            CVulkanContext::Get()->GetMeshVertexInputBindingDescriptions();

        auto VertexInputAttributeDescription=
            CVulkanContext::Get()->GetMeshVertexInputAttributeDescriptions();

        vk::PipelineVertexInputStateCreateInfo PipelineVertexInputStateCreateInfo=vk::PipelineVertexInputStateCreateInfo()
            .setVertexBindingDescriptionCount(VertexInputBindingDescriptions.Num())
            .setPVertexBindingDescriptions(VertexInputBindingDescriptions.Data())
            .setVertexAttributeDescriptionCount(VertexInputAttributeDescription.Num())
            .setPVertexAttributeDescriptions(VertexInputAttributeDescription.Data());

        vk::PipelineInputAssemblyStateCreateInfo PipelineInputAssemblyStateCreateInfo
            =vk::PipelineInputAssemblyStateCreateInfo()
            .setTopology(vk::PrimitiveTopology::eTriangleList);
        
        vk::PipelineViewportStateCreateInfo PipelineViewportStateCreateInfo=vk::PipelineViewportStateCreateInfo()
            .setViewportCount(1)
            .setScissorCount(1);

        vk::PipelineRasterizationStateCreateInfo PipelineRasterizationStateCreateInfo=vk::PipelineRasterizationStateCreateInfo()
            .setPolygonMode(vk::PolygonMode::eFill)
            .setCullMode(vk::CullModeFlagBits::eNone)
            .setFrontFace(vk::FrontFace::eClockwise)
            .setLineWidth(1.0f);

        vk::PipelineMultisampleStateCreateInfo PipelineMultisampleStateCreateInfo=vk::PipelineMultisampleStateCreateInfo()
            .setRasterizationSamples(vk::SampleCountFlagBits::e1);

        vk::PipelineDepthStencilStateCreateInfo PipelineDepthStencilStateCreateInfo=vk::PipelineDepthStencilStateCreateInfo()
            .setDepthTestEnable(true)
            .setDepthWriteEnable(true)
            .setDepthCompareOp(vk::CompareOp::eLessOrEqual)
            .setDepthBoundsTestEnable(false)
            .setStencilTestEnable(false);

        vk::PipelineColorBlendAttachmentState PipelineColorBlendAttachmentState=vk::PipelineColorBlendAttachmentState()
            
                .setColorWriteMask(
                    vk::ColorComponentFlagBits::eR
                    |vk::ColorComponentFlagBits::eG
                    |vk::ColorComponentFlagBits::eB
                    |vk::ColorComponentFlagBits::eA
                )
                .setBlendEnable(true) //color might be transparent
                .setSrcColorBlendFactor(vk::BlendFactor::eSrcAlpha)
                .setDstColorBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
                .setColorBlendOp(vk::BlendOp::eAdd)
                .setSrcAlphaBlendFactor(vk::BlendFactor::eOne)
                .setDstAlphaBlendFactor(vk::BlendFactor::eZero)
                .setAlphaBlendOp(vk::BlendOp::eAdd)
                ;

        vk::PipelineColorBlendStateCreateInfo PipelineColorBlendStateCreateInfo=vk::PipelineColorBlendStateCreateInfo()
            .setAttachmentCount(1)
            .setPAttachments(&PipelineColorBlendAttachmentState);

        vk::GraphicsPipelineCreateInfo GraphicsPipelineCreateInfo=vk::GraphicsPipelineCreateInfo()
            .setPDynamicState(&PipelineDynamicStateCreateInfo)
            .setStageCount(2)
            .setPStages(PipelineShaderStageCreateInfo)
            .setPVertexInputState(&PipelineVertexInputStateCreateInfo)
            .setPInputAssemblyState(&PipelineInputAssemblyStateCreateInfo)
            .setPViewportState(&PipelineViewportStateCreateInfo)
            .setPRasterizationState(&PipelineRasterizationStateCreateInfo)
            .setPMultisampleState(&PipelineMultisampleStateCreateInfo)
            .setPDepthStencilState(&PipelineDepthStencilStateCreateInfo)
            .setPColorBlendState(&PipelineColorBlendStateCreateInfo)
            .setLayout(AxisPipelineLayout)
            .setRenderPass(AxisRenderPass)
            .setSubpass(0);

        AxisPipeline=CVulkanContext::Get()->GetDevice().createGraphicsPipeline(nullptr,GraphicsPipelineCreateInfo).value;
    }   

}

void SForegroundMeshPipeline::UpdateFrameBuffer(vk::Extent2D InSize)
{
    if(AxisColorRenderTarget
        &&AxisColorRenderTarget->GetSize().width==InSize.width
        &&AxisColorRenderTarget->GetSize().height==InSize.height
        )
    {
        return;
    }

    //create render targets
    AxisColorRenderTarget=std::make_shared<CGPURenderTarget>();
    AxisColorRenderTarget->Initialize(
        vk::Extent3D(InSize.width,InSize.height,1)
        , vk::Format::eR8G8B8A8Unorm
        , vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eSampled
    );

    AxisDepthRenderTarget=std::make_shared<CGPURenderTarget>();
    AxisDepthRenderTarget->Initialize(
        vk::Extent3D(InSize.width,InSize.height,1)
        , vk::Format::eD32Sfloat
        , vk::ImageUsageFlagBits::eDepthStencilAttachment
    );

    //create frame buffer
    if( AxisFramebuffer)
    {
        CVulkanContext::Get()->GetDevice().destroyFramebuffer( AxisFramebuffer);
    }

    vk::ImageView Attachments[2]=
    {
         AxisColorRenderTarget->GetImageView(),
         AxisDepthRenderTarget->GetImageView()
    };

    vk::FramebufferCreateInfo FramebufferCreateInfo = vk::FramebufferCreateInfo()
        .setRenderPass(AxisRenderPass)
        .setAttachmentCount(2)
        .setPAttachments(Attachments)
        .setWidth(InSize.width)
        .setHeight(InSize.height)
        .setLayers(1);

     AxisFramebuffer=CVulkanContext::Get()->GetDevice().createFramebuffer(FramebufferCreateInfo);
}

